#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/9/16 14:26
# @Author  : CHEN Wang
# @Site    :
# @File    : fac_contribution.py
# @Software: PyCharm

"""
脚本说明： 计算因子贡献度及因子偏离度，具体算法说明文档见因子择时相关的文档
"""

import pandas as pd
import numpy as np
from quant_researcher.quant.datasource_fetch.stock_api import stock_price_related
from quant_researcher.quant.datasource_fetch.factor_api import factor_exposure_related
from quant_researcher.quant.factors.factor_preprocess.preprocess import my_qcut2
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.datasource_fetch.index_api.index_components import \
    get_index_hist_stock_pool


def fac_contribution_calc(factor_list=None, begin_date=None, end_date=None, external_data=None, window=60,
                          universe='HS300', **kwargs):
    """
    计算因子贡献度或偏离度

    :param list factor_list: list of str, 每个str是股票因子库中的因子英文名称
    :param str begin_date: str, 需要获取因子贡献度的起始日期, 格式如'2020-01-01'
    :param str end_date: str, 需要获取因子贡献度的终止日期， 格式如'2020-05-01'
    :param int window: int, 观察期窗口长度T, 默认60天
    :param str universe: str, 股票池/指数的名称, 如:'HS300', 目前支持
                        {'SH50': '000016', 'SH180': '000010', 'HS300': '399300', 'ZZ500': '000905',
                        'CYBZ': '399006', 'ZXB': '399005', 'ALLA': '000985'}
    :param pd.Dataframe external_data: 外部因子数据，确认数据中至少包含tradedate，stockcode和因子名称三列，多个因子在列上扩充
        一个包含ROE指标的数据形状如下：
        —————————————
        stockcode |tradedate | roe  | quick_ratio
        000001   |2010-01-03| 0.2  | 0.4
       000002   |2010-01-03| 0.15 | 0.2

    :param kwargs:
        - contri_or_devi，str，计算贡献度还是偏离度, 支持'contri' 和 'devi' 默认：'contri'
        - group_num: 分组个数，计算贡献度或者偏离度的分组个数
        - cut_type: 分组类型，quantile：百分位cut， non_quantile: 数值cut
    :return:
    """
    contri_or_devi = kwargs.pop('contri_or_devi', 'contri')
    group_num = kwargs.pop('group_num', 5)
    cut_type = kwargs.pop('cut_type', 'quantile')
    min_group_num = 1
    max_group_num = group_num
    data_begin = max(external_data['tradedate'].min(),
                     begin_date) if external_data is not None else begin_date
    data_end = min(external_data['tradedate'].max(),
                   end_date) if external_data is not None else end_date
    universe_dict = {'SH50': '000016', 'SH180': '000010', 'HS300': '399300', 'ZZ500': '000905',
                     'CYBZ': '399006', 'ZXB': '399005', 'ALLA': '000985'}
    # 因为有观察期的长度 所以开始日期要往前取一个观察期长度的数据才可以算完begin_date和end_date之间的贡献度
    begin = time_tool.date_shifter(data_begin, 'days', -window - 5)
    end = data_end
    asset_pool = get_index_hist_stock_pool(universe_dict[universe], data_begin, data_end)[
        'stock_code'].tolist()
    if external_data is None:
        assert universe in universe_dict.keys(), f"暂不支持该资产池，目前支持的资产池有{universe_dict.keys()}"
    else:
        factor_list = external_data.columns.difference(['tradedate', 'stockcode']).tolist()
        asset_pool = external_data['stockcode'][external_data['stockcode'].isin(asset_pool)].unique().tolist()
    asset_data = stock_price_related.get_stock_quote(asset_pool, begin, end,
                                                     select=['stock_code', 'end_date',
                                                                'high', 'low', 'close as fap_close'])
    asset_data = asset_data.rename(columns={'stock_code': 'stockcode', 'end_date': 'tradedate'})
    asset_data = asset_data.set_index(['tradedate', 'stockcode'])

    # 股票池股票总数
    N = len(asset_pool)
    # Top和Bottom组合中 分别包含的股票数量为 股票池股票总数的20%取整数
    n = int(N / group_num)
    # 过去T天股票的对数收益率
    close_price = asset_data['fap_close'].unstack()
    close_price.index = pd.to_datetime(close_price.index)
    stock_rets_window = np.log(close_price.pct_change(freq=f'{window}D').dropna(how='all') + 1)
    ret_group_mask = stock_rets_window.apply(my_qcut2, group_num=group_num, axis=1, cut_type=cut_type)

    if contri_or_devi == 'contri':
        # 过去T天收益率最低20%的股票的平均收益率
        min_by_ret = stock_rets_window[ret_group_mask==min_group_num].mean(axis=1)
        # min_by_ret = stock_rets_window[stock_rets_window.rank(axis=1).apply(lambda x: x <= n)].mean(
        #     axis=1)
        # 过去T天收益率最高20%的股票的平均收益率
        max_by_ret = stock_rets_window[ret_group_mask==max_group_num].mean(axis=1)

    fac_con_list = []
    fac_dev = pd.DataFrame(index=stock_rets_window.index)
    for factor in factor_list:
        if external_data is None:
            # factor_data = factor_exposure_related.get_stock_factor_exposure(factor, asset_pool, begin, end)
            factor_data = factor_exposure_related.get_stock_factor_exposure(factor, asset_pool, begin, end)
            factor_data = factor_data.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
            factor_data = factor_data.set_index(['tradedate', 'stockcode'])[factor].unstack()
        else:
            factor_data = external_data[external_data.stockcode.isin(asset_pool)]
            factor_data = factor_data.set_index(['tradedate', 'stockcode'])[factor].unstack()
        # 对于因子F，将N支股票依据t-T时刻的因子值由小到大进行排名
        factor_data.index = pd.to_datetime(factor_data.index)
        fac = factor_data.shift(freq=f'{window}D').dropna(how='all')
        fac_group_mask = fac.apply(my_qcut2, group_num=group_num, axis=1, cut_type=cut_type)

        if contri_or_devi == 'contri':

            # 因子值最小的20%的股票的平均收益率
            min_by_fac = stock_rets_window[fac_group_mask==min_group_num].mean(axis=1)
            # 因子值最大的20%的股票的平均收益率
            max_by_fac = stock_rets_window[fac_group_mask==max_group_num].mean(axis=1)
            # 这里计算的因子贡献度和国信研报上算的互为相反数
            # 因子贡献度为正 表示因子值较大的股票组合平均收益较高； 为负表示因子值较大的股票组合平均收益较低
            fac_con_partial = ((max_by_fac - min_by_fac) / (max_by_ret - min_by_ret))
            fac_con_partial = fac_con_partial.rename(factor).loc[begin_date:]
            fac_con_list.append(fac_con_partial)

        if contri_or_devi == 'devi':
            # 对于因子F，将N支股票依据t-T时刻的因子值由小到大进行排名
            # 收益率最低的20%股票的因子排名的均值
            min_rank_by_ret = fac.rank(axis=1)[ret_group_mask==min_group_num].mean(axis=1)
            # 收益率最高的20%股票的因子排名的均值
            max_rank_by_ret = fac.rank(axis=1)[ret_group_mask==max_group_num].mean(axis=1)
            # 这里计算的因子偏离度也和国信研报上算的互为相反数
            # 因子偏离度为正 表示高收益组合中的股票关于因子F的排名分布集中在较高区域； 为负表示关于因子F的排名分布集中在较低区域
            fac_dev_partial = ((max_rank_by_ret - min_rank_by_ret) / (N - n)).rename(factor)

    if contri_or_devi == 'contri':
        fac_con = pd.concat(fac_con_list, axis=1)
        return fac_con
    if contri_or_devi == 'devi':
        fac_dev = fac_dev.join(fac_dev_partial)
        return fac_dev


if __name__ == '__main__':
    CON = db_conn.get_derivative_data_conn()
    CON_factor = db_conn.get_tk_factors_conn()
    factor_list = ['cetop', 'quick_ratio']
    # 需要获取因子贡献度的起始和终止
    begin_date = '2020-01-01'
    end_date = '2020-05-01'
    # 观察期窗口长度T
    window = 60
    universe = 'HS300'

    # 测试内部因子
    fac_con = fac_contribution_calc(factor_list, begin_date, end_date, window=window,
                                    universe=universe, contri_or_devi='contri')
    fac_dev = fac_contribution_calc(factor_list, begin_date, end_date, window=window,
                                    universe=universe, contri_or_devi='devi')

    # 测试外部因子输入的情况
    temp_list = []
    asset_pool = get_index_hist_stock_pool('399300', begin_date, end_date)['stock_code'].tolist()
    for factor in factor_list:
        # temp = factor_exposure_related.get_stock_factor_exposure(factor, asset_pool, begin_date, end_date)
        temp = factor_exposure_related.get_stock_factor_exposure(factor, asset_pool, begin_date, end_date)
        temp = temp.rename(columns={'end_date': 'tradedate', 'stock_code': 'stockcode'})
        temp_list.append(temp.set_index(['tradedate', 'stockcode']))
    external_data = pd.concat(temp_list, axis=1).reset_index()
    external_data.tradedate = external_data.tradedate.astype(str)
    fac_con_ex = fac_contribution_calc(factor_list, begin_date, end_date, window=window,
                                       universe=universe, external_data=external_data,
                                       contri_or_devi='contri')
    fac_dev_ex = fac_contribution_calc(factor_list, begin_date, end_date, window=window,
                                       universe=universe, external_data=external_data,
                                       contri_or_devi='devi')
    CON.close()
    CON_factor.close()
